#include "reconstructionHelpers.hpp"
#include "eigenHelpers.hpp"
#include "unifiedMathFunctions.hpp"
#include "sweFunctions.hpp"

namespace hms::hllc {

/** @brief Returns the normal flow velocity in the star region
 * 
 * @param vNormalL Velocity normal to the edge, left side.
 * @param vNormalR Velocity normal to the edge, right side.
 * @param cwL The wavespeed c_WL, not the HLLC one.
 * @param cwR The wavespeed c_WR, not the HLLC one.
 * 
 * Implements Simons (2019) Eq. 3.57 */
template<typename V, typename C>
auto vNormalStar(
	const V& vNormalL,
	const V& vNormalR,
	const C& cwL,
	const C& cwR
) -> decltype(auto)
{
	assert( notNegative(cwL) && notNegative(cwR) );
	return (vNormalL + vNormalR) / 2. + cwL - cwR;
}

/** @brief Returns the wave speed in the star region.
 * 
 * Simons (2019), Eq. 3.56 gives the water depth in the star region.
 * However, it's never needed by itself, so using this instead skips a couple
 * of calculations.
 */
template<typename V, typename C>
auto cwStar(
	const V& vNormalL,
	const V& vNormalR,
	const C& cwL,
	const C& cwR
) -> decltype(auto)
{
	assert( notNegative(cwL) && notNegative(cwR) );
	return abs( (cwL+cwR)/2. + (vNormalL-vNormalR)/4. );
}


/** @brief Returns the "pressure flux" for a given water depth.
 * 
 * This is equivalent to the hydrostatic force divided by the density, i.e.
 * 1/2 * g * d^2
 * It is part of equations 3.63 and 3.65 in Simons (2019).
 */
template<typename D>
auto pressureFlux( const D& d ) -> decltype(auto)
{
	assert( notNegative(d) );
	return phys::g * square(d) / 2;
}

/**
 * @brief Returns the mass flux d * v_normal as part of Simons (2019), Eq. 3.63
 * 
 * @param d 1xN water depth of the current side
 * @param vNormal 1xN velocity normal to the edge(s)
 * @return decltype(auto) 1xN expression for the mass flux
 */
template<typename D, typename V>
auto massFlux(
	const D& d,
	const V& vNormal
) -> decltype(auto)
{
	assert( notNegative(d) );
	return d * vNormal;
}

/**
 * @brief Returns the first momentum flux component for orthogonal grids,
 * i.e. u^2 * d + 1/2 g * d^2 (Simons (2018), Eq. 3.51, second component).
 * The part "vNormal^2 * d + 1/2 g * d^2" is still part of the momentum flux
 * on non-orthogonal grids.
 * 
 * @tparam A 
 * @tparam B 
 * @param d
 * @param vNormal 
 * @return auto 
 */
template<typename D, typename V>
auto momentumFlux_normal(
	const D& d,
	const V& vNormal
) -> decltype(auto)
{
	return vNormal * massFlux(d, vNormal) + pressureFlux(d);
}


/** @brief Returns the HLL flux. Is used to compute the flux in the star region.
 * Implements Simons (2019), Eq. 3.53:
 * f* = SR * fL - SL*fR + SL*SR*(qR-qL) / (SR - SL)
 * 
 * @param sSide  HLLC wavespeed on the flux' side.
 * @param sOther HLLC wavespeed on the other side.
 * @param fSide  Flux component on the flux' side, outside the star region.
 * @param fOther Flux component on the other side.
 * @param qSide  Storage component on the flux' side.
 * @param qOther Storage component on the other side.
 * @return scalar 
 * 
 */
template<
	typename S1, typename S2,
	typename F, typename Q>
auto hllFlux(
	const S1& sSide,
	const S2& sOther,
	const F& fSide,
	const F& fOther,
	const Q& qSide,
	const Q& qOther
) -> decltype(auto)
{
	/* less binary bloat than a string */
	#define DIVISION_BY_ZERO 1
	assert( nonZero(sSide - sOther) && DIVISION_BY_ZERO );
	#undef DIVISION_BY_ZERO

	#ifndef NDEBUG
	if constexpr ( isDenseType<S1>() ){
		static_assert( isDenseType<S2>() );
		Index nCols { sSide.cols() };
		assert( nCols == sOther.cols() );
		assert( nCols == fSide.cols() );
		assert( nCols == fOther.cols() );
		assert( nCols == qSide.cols() );
		assert( nCols == qOther.cols() );
	}
	if constexpr ( isDenseType<F>() ){
		static_assert( isDenseType<Q>() );
		static_assert( F::RowsAtCompileTime == Q::RowsAtCompileTime );
	}
	#endif

	static constexpr Index nRows { [&]() -> Index {
		if constexpr ( isDenseType<F>() )
			return F::RowsAtCompileTime;
		return 1;
	}() };

	return (
		  repl<nRows>(sOther) * fSide
		- repl<nRows>(sSide)  * fOther
		+ repl<nRows>(sSide * sOther) * (qOther - qSide)
	) / repl<nRows>(sOther - sSide);
}

/**
 * @brief Sets the passed in flux argument to the appropriate values
 * determined by the Riemann scheme.
 * 
 * @tparam isOrtho: On orthogonal meshes, the calculation can be simplified.
 * Then, the orientation template parameter must be provided.
 * @tparam orien Can be either EdgeRange::horizontal or EdgeRange::vertical.
 * Must be provided for orthogonal meshes. If unknown, set isOrtho = false.
 * @param fluxLeft is the out-parameter to which the flux is written.
 * @param sL is the left wave speed
 * @param sR is the right wave speed
 * @param dL is the reconstructed left water depth
 * @param dR is the reconstructed right water depth
 * @param vnL is the reconstructed left velocity normal to the edge
 * @param vnR is the reconstructed right velocity normal to the edge
 * @param vpL is the reconstructed left velocity parallel to the edge
 * @param vpR is the reconstructed right velocity parallel to the edge
 * @param normal is the normal vector. Is ignored if isOrtho=true.
 */
template<bool isOrtho, EdgeRange orien>
HMS_INLINE
void setFlux(
	Ref<Array3s> fluxLeft,
	scalar sL, scalar sR,
	scalar dL, scalar dR,
	scalar vnL, scalar vnR,
	scalar vpL, scalar vpR,
	[[maybe_unused]] const Ref<const Vector2s>& normal
){
	scalar
		vp, f1, f2,
		massFluxL { massFlux(dL, vnL) },
		massFluxR { massFlux(dR, vnR) },
		momFluxL  { momentumFlux_normal(dL, vnL) },
		momFluxR  { momentumFlux_normal(dR, vnR) };

	if ( sL < 0 && sR > 0){
		scalar
			helpL { dL * (vnL - sL) },
			helpR { dR * (vnR - sR) };

		if ( sign(sL * helpR - sR * helpL) == sign(helpR - helpL) ){
			vp = vpL;
			f1 = hllFlux(sL, sR, massFluxL, massFluxR, dL, dR);
			f2 = hllFlux(sL, sR, momFluxL, momFluxR, massFluxL, massFluxR);
		} else {
			vp = vpR;
			f1 = hllFlux(sR, sL, massFluxR, massFluxL, dR, dL);
			f2 = hllFlux(sR, sL, momFluxR, momFluxL, massFluxR, massFluxL);
		}
	} else if (sL >= 0){
		vp = vpL;
		f1 = massFluxL;
		f2 = momFluxL;
	} else {
		vp = vpR;
		f1 = massFluxR;
		f2 = momFluxR;
	}

	if constexpr (isOrtho && orien != EdgeRange::undefined){
		if constexpr ( orien == EdgeRange::horizontal ){
			fluxLeft = Array3s{ f1, -f1*vp, f2 };
		} else {
			fluxLeft = Array3s{ f1, f2, f1*vp};
		}
	} else {
		fluxLeft = Array3s{
			f1,
			normal[0] * f2 - normal[1] * f1 * vp,
			normal[1] * f2 + normal[0] * f1 * vp
		};
	}
}

template<bool isOrtho, EdgeRange orien>
HMS_INLINE
void setFlux(
	Ref<Array3Xs> fluxLeft,
	const Ref<const Array1Xs>& sL,
	const Ref<const Array1Xs>& sR,
	const Ref<const Array1Xs>& dL,
	const Ref<const Array1Xs>& dR,
	const Ref<const Array1Xs>& vnL,
	const Ref<const Array1Xs>& vnR,
	const Ref<const Array1Xs>& vpL,
	const Ref<const Array1Xs>& vpR,
	[[maybe_unused]] const Ref<const Matrix2Xs>& normal
){
	#ifndef NDEBUG
	auto colsAssert = [&]( const auto& a ){
		assert( a.cols() == fluxLeft.cols() );
	};
	colsAssert(sL);
	colsAssert(sR);
	colsAssert(dL);
	colsAssert(dR);
	colsAssert(vnL);
	colsAssert(vnR);
	colsAssert(vpL);
	colsAssert(vpR);
	colsAssert(normal);
	#endif

	for (Index i=0; i<fluxLeft.cols(); ++i){
		setFlux<isOrtho,orien>(
			fluxLeft.col(i),
			sL[i], sR[i],
			dL[i], dR[i],
			vnL[i], vnR[i],
			vpL[i], vpR[i],
			normal.col(i)
		);
	}
}

template<typename SolverType, EdgeRange orien,
	typename F, typename Buf,
	typename D, typename VN, typename VP, typename N
	>
HMS_INLINE void hllcScheme(
	F& flux, Eigen::DenseBase<Buf>& buf,
	const D& dL, const D& dR,
	const VN& vnL, const VN& vnR,
	const VP& vpL, const VP& vpR,
	const N& normal
){
	using S = SolverType;
	static constexpr bool isOrtho {
		std::is_same_v<typename S::MeshType, UniMesh> ||
		std::is_same_v<typename S::MeshType, RectMesh>
	};

	using MapOrScalar =
		std::conditional_t<isDenseType<D>(), Array1XsMap, scalar>;

	auto mapToCol = [&]([[maybe_unused]] Index i) -> MapOrScalar {
		if constexpr ( isDenseType<D>() ){
			const Index& nEdges { dL.cols() };
			assert( buf.cols() >= 6 );
			assert( buf.rows() >= nEdges );
			return { buf.col(i).data(), nEdges };
		} else {
			return {};
		}
	};
	/* wave speeds */
	MapOrScalar
		cwL    { mapToCol(0) },
		cwR    { mapToCol(1) },
		vStar  { mapToCol(2) },
		cwStar { mapToCol(3) },
		sL     { mapToCol(4) },
		sR     { mapToCol(5) };

	cwL = hms::wavespeed(dL);
	cwR = hms::wavespeed(dR);

	/* water depth and wave speed in the star region */
	vStar  =  vNormalStar(vnL, vnR, cwL, cwR);
	cwStar = hllc::cwStar(vnL, vnR, cwL, cwR);

	sL = select(dL == 0,
		vnR - 2 * cwR,
		min(vnL - cwL, vStar - cwStar)
	);
	sR = select(dR == 0,
		vnL + 2 * cwL,
		max(vnR + cwR, vStar + cwStar)
	);

	using hms::reconstruction::left;
	using hms::reconstruction::right;

	setFlux<isOrtho, orien>(
		left(flux), sL, sR, dL, dR, vnL, vnR, vpL, vpR, normal
	);
	/* to be able to unify functions between batch wise and single edge
	 * approach, we need to provide the right type of zero */
	auto zero = [&](){
		if constexpr( isDenseType<D>() )
			return scalar{0};
		else
			return Array3s::Zero().eval();
	};

	static constexpr auto dryDepth { S::minWaterDepth };
	left(flux) = select( repl<3>( dL < dryDepth && dR < dryDepth ),
		zero(), left(flux)
	);
	right(flux) = left(flux);
}

template<typename SolverType, EdgeRange orien>
ArrayNNs<3,2> hllcScheme(
	const EdgeValues& ev,
	const ArrayNNs<2,2>& vRotLR,
	const Ref<const Vector2s>& normal
){
	ArrayNNs<3,2> flux;
	Array1s dummyBuf;
	hllcScheme<SolverType, orien>(
		flux, dummyBuf,
		ev.dLR[0], ev.dLR[1],
		vRotLR(0,0), vRotLR(0,1), vRotLR(1,0), vRotLR(1,1),
		normal
	);
	return flux;
}

} // end namespace hms::hllc